#include "precond.hpp"

/// @brief ILU预条件子
/// @tparam idx_t 索引类型
/// @tparam data_t 存储数据的浮点类型
/// @tparam calc_t 执行计算时的浮点类型
/// @tparam setup_t setup所用的浮点类型
template<typename idx_t, typename data_t, typename calc_t, typename setup_t>
class ILU: public Precond<idx_t, calc_t, setup_t>
{
private:
    // 本进程负责的行数
    idx_t nrows = 0, addS_nrows = 0;
    // 调用者（外部）的矩阵
    MyParCSRMatrix<idx_t, setup_t, setup_t>* A_calc = nullptr;
    
    std::vector<idx_t>  L_rpt,// 分解后的上三角矩阵的行偏移
                        L_cid;// 分解后的上三角矩阵的列索引
    
    std::vector<idx_t>  U_rpt,// 分解后的下三角矩阵的行偏移
                        U_cid;// 分解后的下三角矩阵的列索引
    
    std::vector<data_t> L_val,// 分解后的上三角矩阵的值
                        U_val;// 分解后的下三角矩阵的值
    calc_t * aux = nullptr;
public:
    ILU() {}
    ~ILU() {
        if (aux) { if (aux) { delete [] aux; aux = nullptr; } }
    }
    /// @brief ILU预条件的建立函数
    /// @param par_A 外部调用者的矩阵A
    void Setup(MyParCSRMatrix<idx_t, setup_t, setup_t> & par_A)
    {
        A_calc = & par_A;
        nrows = par_A.diag->nrows; assert(nrows == par_A.end_row - par_A.beg_row);
        /*
        addS_nrows = par_A.diag->nrows + par_A.offd->ncols;
        // 首先组装带有Schwartz的矩阵
        
        for (idx_t s = 0; s < par_A.num_sends; s++) {
            idx_t       * buf      = (idx_t *) (par_A.send_buf)     + par_A.send_cnts[s];
            const idx_t * glb_rows = par_A.send_row_ids + par_A.send_cnts[s];
            const idx_t send_pid = par_A.send_pids[s];
            const idx_t num = send_cnts[s + 1] - send_cnts[s];
            for (idx_t k = 0; k < num; k++) {
                buf[k] = par_A.
            }
        }
        */
        
        aux = new calc_t [nrows];
        // 分配行偏移的大小
        L_rpt.resize(nrows + 1, 0);
        U_rpt.resize(nrows + 1, 0);

        const MyCSRMatrix<idx_t, setup_t, setup_t> & A = *(par_A.diag);
        // 执行ILU分解过程
        for (idx_t i = 0; i < nrows; i++) {// 逐行遍历原来的A矩阵
            const idx_t nnz = A.row_ptr[i + 1] - A.row_ptr[i];// 第i行的非零元数目
            const idx_t offset = A.row_ptr[i];// 第i行的非零元在A矩阵中的起始偏移
            setup_t rbuf[nnz];// 拷入当前行的值
            idx_t lnz = 0;// 记录该行有多少个左非零元（位于对角线左侧）
            for (idx_t p = 0; p < nnz; p++) {// 将该行的非零元拷入缓冲区准备计算
                rbuf[p] = A.vals[p + offset];
                if (i == A.col_idx[p + offset]) lnz = p;
            }
            idx_t rnz = nnz - lnz;// 记录该行有多少个右非零元（包含对角元）
            // 逐个消去左边的元素
            for (idx_t lp = 0; lp < lnz; lp++) {
                const idx_t lj = A.col_idx[lp + offset];// 欲消去的该非零元的列号
                const data_t * u_vals = & U_val[ U_rpt[lj] ];// 该列对应的那一行的非零元值（注意是之前已经分解过的U了）
                const idx_t  * u_cids = & U_cid[ U_rpt[lj] ];// 该列对应的那一行的列索引（注意是之前已经分解过的U了）
                assert(u_cids[0] == lj);// 
                rbuf[lp] /= u_vals[0];// 首位总是对角元
                // 欲消该元会对该元右边元素产生影响
                for (idx_t rp = lp + 1; rp < nnz; rp++) {// 逐个遍历右非零元，将该影响累加上去
                    const idx_t rj = A.col_idx[rp + offset];
                    // 检查一下矩阵的第lj行是否有列为rj的非零元
                    for (idx_t k = 0; k < U_rpt[lj + 1] - U_rpt[lj]; k++) {
                        if (u_cids[k] == rj) {// 如果有，就会产生累加
                            rbuf[rp] -= rbuf[lp] * u_vals[k];
                        }
                    }// k loop
                }// rp loop
            }// lp loop
            // 将缓冲区的数据拷回到L和U中
            for (idx_t lp = 0; lp < lnz; lp++) {
                L_cid.emplace_back(A.col_idx[lp + offset]);
                L_val.emplace_back(rbuf     [lp         ]);
            }
            for (idx_t up = lnz; up < nnz; up++) {
                U_cid.emplace_back(A.col_idx[up + offset]);
                U_val.emplace_back(rbuf     [up         ]);
            }
            // 记录行偏移
            L_rpt[i + 1] = L_rpt[i] + lnz; assert(L_rpt[i + 1] == (idx_t) L_cid.size());
            U_rpt[i + 1] = U_rpt[i] + rnz; assert(U_rpt[i + 1] == (idx_t) U_cid.size());
        }
    }
    /// @brief 前向求解：求解下三角方程组 L*sol = rhs
    /// @param rhs 右端项
    /// @param sol 解向量
    void forward(const calc_t * rhs, calc_t * sol) {
        for (idx_t i = 0; i < nrows; i++) {
            calc_t tmp = rhs[i];
            for (idx_t p = L_rpt[i]; p < L_rpt[i + 1]; p++) {
                assert(L_cid[p] < i);
                tmp -= L_val[p] * sol[L_cid[p]];
            }
            sol[i] = tmp;
        }
    }
    /// @brief 后向求解：求解上三角方程组 U*sol = rhs
    /// @param rhs 右端项
    /// @param sol 解向量
    void backward(const calc_t * rhs, calc_t * sol) {
        for (idx_t i = nrows - 1; i >= 0; i--) {
            calc_t tmp = rhs[i];
            const idx_t offset = U_rpt[i]; 
            const idx_t * u_cids = & U_cid[offset];
            const data_t* u_vals = & U_val[offset];
            assert(u_cids[0] == i);
            const idx_t num = U_rpt[i + 1] - offset;
            for (idx_t p = 1; p < num; p++) {
                assert(u_cids[p] > i);
                tmp -= u_vals[p] * sol[u_cids[p]];
            }
            sol[i] = tmp / u_vals[0];
        }
    }
    /// @brief ILU预条件的求解函数
    /// @param b 右端项
    /// @param x 解向量
    /// @param zero_guess 是否使用零初值优化 
    void Solve(const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & x, bool zero_guess=true)
    {
        assert(b.end_row - b.beg_row == nrows);
        assert(b.end_row == x.end_row && b.beg_row == x.beg_row);

        if (zero_guess) {// 使用零初值优化时, L*U*x = b
            forward (b.data, aux);// L*aux = b，令aux = U*x
            backward(aux, x.data);// U*x = aux
        }
        else {// 非零初值时，需要先计算一遍当前的残差
            MyParCSRVector<idx_t, calc_t> tmp(b);
            A_calc->Mult((calc_t) -1.0, x, (calc_t) 1.0, b, tmp);// tmp <- b - A*x (i.e., residual)
            // 再对残差使用ILU分解的前代回代
            forward (tmp.data, aux);// L*aux = r
            backward(aux, tmp.data);// U*tmp = aux
            vec_add(x, (calc_t) 1.0, tmp, x);// x += tmp
        }
    }
};